#include <iostream>
#include <cstring>
#include <sys/socket.h>
#include <netinet/in.h>
#include <netinet/ip_icmp.h>
#include <arpa/inet.h>
#include <unistd.h>
#include <sys/time.h>

#define PACKET_SIZE 64
#define MAX_WAIT_TIME 5

unsigned short checksum(void *b, int len) {
    unsigned short *buf = (unsigned short *)b;
    unsigned int sum = 0;
    unsigned short result;

    for (sum = 0; len > 1; len -= 2) {
        sum += *buf++;
    }
    if (len == 1) {
        sum += *(unsigned char *)buf;
    }
    sum = (sum >> 16) + (sum & 0xFFFF);
    sum += (sum >> 16);
    result = ~sum;
    return result;
}

int main(int argc, char *argv[]) {
    if (argc != 2) {
        std::cerr << "Usage: " << argv[0] << " <server_ip>" << std::endl;
        return 1;
    }

    int sockfd;
    struct sockaddr_in serv_addr;
    char packet[PACKET_SIZE];
    char recv_packet[PACKET_SIZE];
    
    if ((sockfd = socket(AF_INET, SOCK_RAW, IPPROTO_ICMP)) < 0) {
        perror("socket");
        return 1;
    }

    memset(&serv_addr, 0, sizeof(serv_addr));
    serv_addr.sin_family = AF_INET;
    if (inet_pton(AF_INET, argv[1], &serv_addr.sin_addr) <= 0) {
        perror("inet_pton");
        close(sockfd);
        return 1;
    }

    struct icmphdr *icmph = (struct icmphdr *)packet;
    icmph->type = ICMP_ECHO;
    icmph->code = 0;
    icmph->un.echo.id = getpid();
    icmph->un.echo.sequence = 0;
    memset(packet + sizeof(struct icmphdr), 0, PACKET_SIZE - sizeof(struct icmphdr));
    icmph->checksum = checksum((unsigned short *)icmph, sizeof(struct icmphdr));

    struct timeval start, end;
    gettimeofday(&start, NULL);

    if (sendto(sockfd, packet, sizeof(struct icmphdr), 0, 
              (struct sockaddr *)&serv_addr, sizeof(serv_addr)) <= 0) {
        perror("sendto");
        close(sockfd);
        return 1;
    }

    std::cout << "Requesting time from server..." << std::endl;

    fd_set readfds;
    FD_ZERO(&readfds);
    FD_SET(sockfd, &readfds);
    struct timeval tv;
    tv.tv_sec = MAX_WAIT_TIME;
    tv.tv_usec = 0;

    int ret = select(sockfd + 1, &readfds, NULL, NULL, &tv);
    if (ret == 0) {
        std::cerr << "Request timed out." << std::endl;
        close(sockfd);
        return 1;
    } else if (ret < 0) {
        perror("select");
        close(sockfd);
        return 1;
    }

    socklen_t addr_len = sizeof(serv_addr);
    ssize_t n = recvfrom(sockfd, recv_packet, PACKET_SIZE, 0, 
                        (struct sockaddr *)&serv_addr, &addr_len);
    if (n < 0) {
        perror("recvfrom");
        close(sockfd);
        return 1;
    }

    gettimeofday(&end, NULL);
    double rtt = (end.tv_sec - start.tv_sec) * 1000.0 + 
                (end.tv_usec - start.tv_usec) / 1000.0;

    struct iphdr *iph = (struct iphdr *)recv_packet;
    int iphdrlen = iph->ihl * 4;

    struct icmphdr *recv_icmph = (struct icmphdr *)(recv_packet + iphdrlen);

    if (recv_icmph->type == ICMP_ECHOREPLY) {
        char *time_data = recv_packet + iphdrlen + sizeof(struct icmphdr);
        std::cout << "Server time: " << time_data << std::endl;
        std::cout << "Round-trip time: " << rtt << " ms" << std::endl;
    } else {
        std::cerr << "Received non-echo-reply packet." << std::endl;
    }

    close(sockfd);
    return 0;
}